#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : chat_model 
# @File    : dialogue_logic.py
# @IDE     : PyCharm 
# @Author  :ZH
# @Time    : 2025/1/3 18:19
__all__ = ['dialogue_logic']

import datetime
from typing import Union

from apps import cruds
from apps.api.v1.chat_api.schema import ChatDialogueListQuery
from apps.api.v1.draw_api.schema import DrawDialogueListQuery
from apps.models import User


class DialogueLogic:
    def __init__(self):
        self.crud = cruds.dialogue_crud

    async def dialogue_list(self, query: Union[ChatDialogueListQuery, DrawDialogueListQuery], user: User):
        query_sql = self.crud.model.filter(session_id=query.session_id, user_id=user.id, is_delete=False)
        count = await query_sql.count()
        if count:
            data = await query_sql.order_by('-id').offset((query.page - 1) * query.limit).limit(query.limit).all()
        else:
            data = []
        return count, data

    async def create_dialogue(self, user: User, session_id: int, query: str, answer: str, model_id, model_name,
                              tool: list = None, usage_tokens: int = 0):
        data = await self.crud.create(obj_in={
            'session_id': session_id,
            'user_id': user.id,
            'query': query,
            'answer': answer,
            'model_id': model_id,
            'model_name': model_name,
            "tool": tool if tool else [],
            "usage_tokens": usage_tokens
        })
        return data

    async def clear_dialogue(self, user: User, session_id: int):
        return await self.crud.model.filter(
            user_id=user.id, session_id=session_id, is_delete=False
        ).update(is_delete=True)

    async def get_default_model(self):
        return await self.crud.get_one(is_default=True, is_delete=False)

    async def home_dialogue(self):
        # 查询最近一个月记录
        return await self.crud.model.filter(
            create_date__gte=datetime.datetime.now() - datetime.timedelta(days=30)
        ).all()


dialogue_logic = DialogueLogic()
